#include "leecode.h"

#include <vector>
#include <algorithm>
#include <string>
#include <array>
#include <set>

struct Task{
  int count = 0;
  char ch;
};

//bool taskCompareLess(const Task& t1, const Task& t2) {
//    return t1.count < t2.count || ((t1.count == t2.count) && (t1.ch < t2.ch));
//}

struct taskCompareLess {
  bool operator()(const Task& t1, const Task& t2) const{
    return t1.count < t2.count || ((t1.count == t2.count) && (t1.ch < t2.ch));
  }
};


class Solution {
public:
    int leastInterval(vector<char>& tasks, int n) {
      initSet(tasks);
      std::vector<char> sequence;
      int seq_len = 0;
      while(!tasks.empty()) {
        bool task_inserted = insertTask(sequence, n);
        if(!task_inserted) {
          sequence.push_back('0');
        }
      }

      return sequence.size();

    }

    Solution()  {

    }

    bool insertTask(std::vector<char>& sequence, int n ) {
      int seq_len = sequence.size();
      if(seq_len == 0) {
        auto iter = task_remains.begin();
        sequence.push_back(iter->ch);
        Task task{iter->count - 1, iter->ch};
        task_remains.erase(iter);
        if(task.count > 0)
          task_remains.insert(task);
        return true;
      } else {
        for(auto iter = task_remains.begin(); iter != task_remains.end(); iter++) {
          bool has_equal = false;
          for(int i = seq_len - 1; (i >= 0 && i >= seq_len - n); i--) {
            if(sequence[i] == iter->ch) {
              has_equal = true;
              break;
            }
          }
          if(!has_equal) {
            sequence.push_back(iter->ch);
            Task task{iter->count - 1, iter->ch};
            task_remains.erase(iter);
            if(task.count > 0)
              task_remains.insert(task);
            return true;
            break;
          }
        }
      }
    }

    void initSet(vector<char>& tasks){
      map<char, int> task_remains_map;
      for(int i = 0; i < tasks.size(); i++) {
        auto iter = task_remains_map.find(tasks[i]);
        if(iter == task_remains_map.end())
          task_remains_map[tasks[i] ] = 1;
        else
          task_remains_map[tasks[i] ] ++;
      }

      for(auto& pair: task_remains_map) {
        task_remains.insert({pair.second, pair.first});
      }
    }

    set<Task, taskCompareLess> task_remains;
};

namespace {
  void TestInput621(std::vector<char>& sequence, int n ) {
    Solution s;
    int ret = s.leastInterval(sequence, n);

    cout << "input is: " << sequence.data() << "\t\t" << n << endl;
    cout << "result is: " << ret << endl;
  }
}


TEST(leastInterval, test1) {
  GTEST_SKIP();
  std::vector<char> sequence{'A','A','A','B','B','B'};
  TestInput621(sequence, 2);
}
